import os
from pathlib import Path

# 设置路径
# base_path = r"E:\data\ai_project\public_security\infrared_detection\25年\9月\9月30日\红外人像0930"
# base_path = r"E:\zhihao\2025文档\25年10月\ai_训练\ai项目_备份\ljx\infrared_detection_model\datasets\train"
# base_path = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\HIT-UAV\HIT-UAV\data\train"
# base_path = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\HIT-UAV\HIT-UAV\data\val"
# base_path = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\重清洗\lyh\train"
# base_path = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\重清洗\hx\train"
# base_path = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\重清洗\lyh\valid"
base_path = r"E:\zhihao\2025文档\25年10月\ai_训练\数据集\重清洗\hx\valid"
images_dir = os.path.join(base_path, "images")
labels_dir = os.path.join(base_path, "labels")


# 获取所有文件名（不含扩展名）
def get_file_stems(directory, extensions):
    """获取指定目录下指定扩展名的文件名（不含扩展名）"""
    file_stems = set()
    if os.path.exists(directory):
        for file in os.listdir(directory):
            if any(file.lower().endswith(ext) for ext in extensions):
                stem = os.path.splitext(file)[0]
                file_stems.add(stem)
    return file_stems


# 图片常见扩展名
image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']
# 标签扩展名
label_extensions = ['.txt', '.xml']

# 获取文件名集合
image_stems = get_file_stems(images_dir, image_extensions)
label_stems = get_file_stems(labels_dir, label_extensions)

# 找出不匹配的文件
images_without_labels = image_stems - label_stems  # 有图片但没有标签
labels_without_images = label_stems - image_stems  # 有标签但没有图片

# 输出结果
print("=" * 60)
print("文件匹配检查结果")
print("=" * 60)
print(f"\n图片总数: {len(image_stems)}")
print(f"标签总数: {len(label_stems)}")
print(f"匹配数量: {len(image_stems & label_stems)}")

if images_without_labels:
    print(f"\n❌ 有图片但没有对应标签的文件 ({len(images_without_labels)} 个):")
    print("-" * 60)
    for name in sorted(images_without_labels):
        print(f"  - {name}")
else:
    print("\n✅ 所有图片都有对应的标签")

if labels_without_images:
    print(f"\n❌ 有标签但没有对应图片的文件 ({len(labels_without_images)} 个):")
    print("-" * 60)
    for name in sorted(labels_without_images):
        print(f"  - {name}")
else:
    print("\n✅ 所有标签都有对应的图片")

print("\n" + "=" * 60)

# 可选：将结果保存到文件
save_result = input("\n是否将结果保存到文件? (y/n): ")
if save_result.lower() == 'y':
    output_file = os.path.join(base_path, "mismatch_report.txt")
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write("文件匹配检查报告\n")
        f.write("=" * 60 + "\n")
        f.write(f"图片总数: {len(image_stems)}\n")
        f.write(f"标签总数: {len(label_stems)}\n")
        f.write(f"匹配数量: {len(image_stems & label_stems)}\n\n")

        if images_without_labels:
            f.write(f"有图片但没有标签 ({len(images_without_labels)} 个):\n")
            for name in sorted(images_without_labels):
                f.write(f"  - {name}\n")
            f.write("\n")

        if labels_without_images:
            f.write(f"有标签但没有图片 ({len(labels_without_images)} 个):\n")
            for name in sorted(labels_without_images):
                f.write(f"  - {name}\n")

    print(f"✅ 结果已保存到: {output_file}")